from functionNameRegistry import FunctionNameRegistry
import jeringa
import gdb

# An SCC (Self-Correction Code) Detector which identifies Robusto functions.
# On instantiation it must be passed a list of program exit functions used
# inside Robusto - that is, if Robusto has calls to a program-terminating function
# such as exit(), RobustoDetector must know of it to set the corresponding
# RobustoProgramExitBreakpoints.
# USED BY: InjectionRunDriver.
# USES: RobustoFinishBreakpoint, RobustoProgramExitBreakpoint.

class RobustoDetector(object):
    def __init__(self, robustoProgramExitFunctions=["robustoVotingErr"]):
        self.functionNameRegistry = FunctionNameRegistry(functions={"regex": ["NSRobusto.*"]})
        self.SCCExecutionError = None
        self.finishedExecutingSCC = None
        self.setProgramExitDetectors(robustoProgramExitFunctions)

    def isSCC(self, currentPC, functionName):
        return self.functionNameRegistry.isRegistered(functionName)

    def executeSCC(self, currentPC, functionName):
        self.SCCExecutionError = False
        self.finishedExecutingSCC = False

        # Set a FinishBreakpoint at the end of the current frame.
        robustoFinishBreakpoint = RobustoFinishBreakpoint(self)
        robustoStopBreakpoint = RobustoStopBreakpoint(self)

        # Continue until we exit the current frame.
        while not self.finishedExecutingSCC and not self.SCCExecutionError:
            gdb.execute("step")

        # If the FinishBreakpoint still exists (i.e. we didn't hit it), delete it.
        if robustoFinishBreakpoint.is_valid():
            robustoFinishBreakpoint.delete()

        if robustoStopBreakpoint.is_valid():
            robustoStopBreakpoint.delete()

        return self.SCCExecutionError

    def setSCCExecutionError(self, SCCExecutionError):
        self.SCCExecutionError = SCCExecutionError

    def setFinishedExecutingSCC(self, finishedExecutingSCC):
        self.finishedExecutingSCC = finishedExecutingSCC

    def setProgramExitDetectors(self, robustoProgramExitFunctions):
        for programExitFunction in robustoProgramExitFunctions:
            RobustoProgramExitBreakpoint(programExitFunction, self)

################################################################################

# A subclass of gdb.FinishBreakpoint which is set at the return point of Robusto
# functions by a RobustoDetector. When RobustoDetector's executeSCC steps
# past the Robusto function's end point, this will make it stop advancing
# so we can continue driving the target normally.
# If an injection causes the Robusto function to return to somewhere
# other than where this breakpoint is set at, this breakpoint will
# let the RobustoDetector that set it that an error occurred.
# COMMUNICATES WITH: RobustoDetector.

class RobustoFinishBreakpoint(gdb.FinishBreakpoint):
    def __init__(self, robustoDetector):
        super(RobustoFinishBreakpoint, self).__init__()
        self.robustoDetector = robustoDetector

    def out_of_scope(self):
        self.robustoDetector.setSCCExecutionError(True)
        jeringa.log("Breakpoint %d: ERROR EN ROBUSTO (out_of_scope)!!" % self.number)

class RobustoStopBreakpoint(gdb.Breakpoint):
    def __init__(self, robustoDetector):
        super(RobustoStopBreakpoint, self).__init__("*%d" % gdb.selected_frame().older().pc())
        self.robustoDetector = robustoDetector

    def stop(self):
        self.robustoDetector.setFinishedExecutingSCC(True)
        return True   # Stop at this breakpoint.

################################################################################

# A Breakpoint that's set at any program-terminating functions such as exit()
# that may be called inside Robusto's code. This is necessary so that we
# can quickly detect whether Robusto is trying to terminate the program and
# act in consequence.
# COMMUNICATES WITH: RobustoDetector.

class RobustoProgramExitBreakpoint(gdb.Breakpoint):
    def __init__(self, spec, robustoDetector):
        super(RobustoProgramExitBreakpoint, self).__init__(spec)
        self.robustoDetector = robustoDetector

    def stop(self):
        self.robustoDetector.setSCCExecutionError(True)
        jeringa.log("Breakpoint %d: EXIT EN ROBUSTO!!" % self.number)

        return True   # Stop at this breakpoint.
